fix(kimi-k3): move the media wrapper and thinking-effort default to the prompt-encoding layer - #1995

sequenceDiagram
  participant Request
  participant TiktokenTokenizer
  participant KimiK3XTML
  Request->>TiktokenTokenizer: apply KimiK3Xtml chat template
  TiktokenTokenizer->>KimiK3XTML: call served renderer
  KimiK3XTML->>KimiK3XTML: resolve explicit effort or default to max
  KimiK3XTML-->>TiktokenTokenizer: return rendered XTML

_{Comment @coderabbitai help to get the list of available commands.}

👋 The PR description doesn't fully follow
PULL_REQUEST_TEMPLATE.md:

• Missing header: ## Description
• Missing header: ### Problem
• Missing header: ### Solution
• Missing header: ## Changes
• Missing header: ## Test Plan

Please update the PR description so reviewers have the context they need.

Node validation: prompt shape and token parity

Ran this branch on the B300 node against the real Kimi-K3 checkpoint, using vLLM's own server as the reference rather than a hand-written expectation.

Setup. One vLLM gRPC engine (:8081, /raid/models/moonshotai/Kimi-K3, TP8). Two gateways built from source and run natively against that same engine, so the only variable is the SMG binary:

Reference. Before handing the GPUs to the gRPC engine, I captured the same requests against vLLM's HTTP server (:8200) on the same checkpoint — that path does its own chat templating and its own preprocessing, so it is the behaviour we are trying to match. Asking for prompt_logprobs makes the engine echo back the decoded prompt, which gives a direct read of the wrapper instead of inferring it:

<|media_begin|>image 1024x768<|media_content|><|media_pad|>…<|media_pad|><|media_end|>

for a 1024×768 image — 1036 pads, and 3 marker tokens + 6 tokens of image 1024x768 = the 9-token wrapper.

Token parity

usage.prompt_tokens, same requests to all three:

new matches the reference on 8/8. old matches on 2/8 — exactly the two cases that pass an explicit effort and contain no image, i.e. the two things the old build already got right.

Every baseline gap decomposes exactly, with no remainder:

• text-only, 89 − 22 = 67 → the missing thinking-effort directive.
• 1 image with explicit effort, 1139 − 1130 = 9 → the missing wrapper, directive already present.
• 1 image, default effort, 1138 − 1062 = 76 = 67 + 9.
• 1 image 4000×3000, 15547 − 15470 = 77 = 67 + 10 — the wrapper is 10 here because image 4000x3000 is 7 tokens, not 6.
• 2 images, 16594 − 16508 = 86 = 67 + 9 + 10, and 2183 − 2098 = 85 = 67 + 9 + 9.

The per-image term tracks the actual dimensions rather than being a constant, which is the part that matters.

Gateway debug logging agrees from the inside — for the 4000×3000 case the new build reports item_sizes=[(4000, 3000)], original_len=94 expanded_len=15547 placeholder_count=1 with 15444 pads, so the wrapper's non-pad cost is 15547 − 93 − 15444 = 10.

Width×height ordering, and pre-resize dimensions

Token counts cannot catch a W/H swap: swapping the operands of x leaves the token multiset unchanged, so image 4000x3000 and image 3000x4000 both cost 7. I verified that against the checkpoint tokenizer rather than assuming it, then tested ordering a different way — asking the model to repeat the dimensions line from its own context. The numbers only exist there as text, so a correct answer is a direct read of the wrapper.

Strip images make this unambiguous: the same two numbers in both orders, and few enough pads that retrieval is reliable.

The mirror pair both read back correctly, so this is not the model guessing an orientation. old has no dimensions in context and guesses, as expected. Prompt tokens differ by a constant 9 (497 vs 488, 838 vs 829).

This also settles pre- vs post-resize: a 4000×64 strip is heavily resized for patching (~380 pads), yet the wrapper still states the original 4000×64 — item_sizes is captured before the resize.

One caveat

SMG does not forward prompt_logprobs to the engine, so the byte-exact decoded prompt is only available on the vLLM path. On the SMG path the evidence is token parity, the gateway's own expansion counts, and the read-back above. I did try to falsify the wrapper contents rather than only confirm them: for the 4000×3000 image the model answered image 3465x2599, which would be alarming, but that string also costs 7 tokens so parity could not rule it out. It is a long-context retrieval failure across 15444 identical pad tokens — the same probe reads correctly whenever the pad run is short, and the mirror pair pins the ordering independently.

Vendor-verifier accuracy check

Ran the two image benchmarks from MoonshotAI/Kimi-Vendor-Verifier against this branch, to confirm the new media wrapper doesn't cost anything end-to-end.

The three OCRBench numbers sit inside one standard error of each other, so this reads as no regression rather than an improvement. The middle row is the useful one for this PR specifically: it sends no thinking_effort at all, so the score comes from the max default that this PR installs in the K3 XTML render path.

Setup:

• native smg built from this branch (--reasoning-parser kimi_k3 --tool-call-parser kimi_k3), in front of a TP8 gRPC vLLM engine serving moonshotai/Kimi-K3 on a B300 node
• unmodified KVV harness via eval.py <bench> --thinking --think-mode opensource --temperature 1.0 --top-p 0.95 --max-connections 50 --epochs 1 --stream, --max-tokens 16384 for OCRBench and 98304 for MMMU Pro
• token usage — OCRBench 1,349,304 (I 1,100,238 / O 249,066); MMMU Pro 10,209,534 (I 1,655,538 / O 8,553,996)

Two things to note about the numbers rather than the code:

• The harness pins MMMU Pro to the standard (10 options) subset (it reports the dataset as MMMU_Pro_10c), so that row is the standard 10-option split, not the vision-only split.
• One of the 1000 OCRBench samples hit the 16384-token cap mid-generation and was scored as-is; the harness runs with fail_on_error: False. MMMU Pro finished with no errors.